from DrissionPage import SessionPage
from colorama import Fore, init
import sys

# 创建页面对象
page = SessionPage()


def tree(ele):
    init()
    e = ele
    print(f'{Fore.BLUE}{Fore.CYAN}<{e.tag}>  {Fore.RESET}{e.attrs}')
    __tree(e)


def __tree(ele: any, layer=6, has_next_brother=True, body=''):
    try:
        list_ele = ele.children(timeout=0.1)
    except:
        list_ele = []

    length = len(list_ele)
    body_unit = '│   ' if has_next_brother else '    '
    tail = '├───'
    new_body = body + body_unit

    if length > 0 and layer >= 1:
        has_next_brother2 = True
        for i in range(length):
            if i == length - 1:
                tail = '└───'
                has_next_brother2 = False
            e = list_ele[i]
            all_body = f'{Fore.BLUE}{new_body}{tail}{Fore.RESET}'

            print(f'{all_body}{Fore.CYAN}<{e.tag}>{Fore.RESET} ')
            try:
                tree_attr(e, all_body, has_next_brother2, layer)
            except:
                pass

            __tree(e, layer - 1, has_next_brother2, new_body)


def tree_attr(ele, body, has_next_brother=True, layer=3):
    e: dict = ele.attrs
    has_child = True if ele.child(timeout=0.2) else False
    if layer == 1:
        has_child = False

    part1 = '│' if has_next_brother else ' '
    part2 = '│' if has_child else ' '
    replace_part = part1 + '   ' + part2
    new_body = body.replace('├───', replace_part).replace('└───', replace_part)

    text = ele.text.split('\n')[0]
    if len(text) >= 1:
        e['inner_txt'] = text

    if len(e) > 0:
        # e['xpath'] = ele.xpath
        # e['css_path'] = ele.css_path

        max_k_len = max([len(key) for key in e.keys()])
        for k, v in e.items():
            key = Fore.YELLOW + str(k).ljust(max_k_len) + Fore.RESET

            print(new_body, ' ', key, ':', v)


if __name__ == '__main__':
    url = r'https://news.ifeng.com/'
    arguments = sys.argv

    if len(arguments) > 1:
        if 'http://' in arguments[1] or 'https://' in arguments[1]:
            url = arguments[1]

    page.get(url)

    e = page.ele('x:/html')
    tree(e)
    e.css_path
